
#include "lp_norm_v2_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cmath>
#include <limits>

#define EPSILON 1e-6
namespace optiling {
const uint32_t BLOCK_SIZE = 64;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  LpNormV2TilingData tiling;
  uint32_t sizeofdatatype;
  uint32_t totalLengthAligned;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto socVersion = ascendcPlatform.GetSocVersion();
  uint64_t ub_size;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  auto aivNum = ascendcPlatform.GetCoreNumAiv();  // vector core num  1
  
  float p = *(context->GetAttrs()->GetFloat(0));
  float epsilon = std::numeric_limits<float>::epsilon();
  uint32_t pType = 0;
  tiling.set_pValue(p);
  if (std::fabs(p - 2.0f) < epsilon){
    tiling.set_pType(0);
  }else if (std::isinf(p)) {
    if (p > 0) {
      tiling.set_pType(1);
      pType = 1;
    } else {
      tiling.set_pType(2);
      pType = 2;
    }
  } else if (std::fabs(p) < epsilon) {
    tiling.set_pType(3);
    pType = 3;
  } else if (std::fabs(p - 1.0f) < epsilon) {
    tiling.set_pType(4);
    pType = 4;
  } else {
    tiling.set_pType(5);
    pType = 5;
  }

  const gert::TypedContinuousVector<int64_t> *axes = context->GetAttrs()->GetListInt(1);
  
  uint32_t totalLength = 1;//context->GetInputTensor(0)->GetShapeSize();
  uint32_t stepSize = 1;
  uint32_t unitCount = 1;
  auto dim_1 = axes->GetSize();
  if(dim_1 != 0)
  {
    auto shape = context->GetInputTensor(0)->GetShape().GetOriginShape();
    uint32_t dim_num = shape.GetDimNum();
    uint32_t maxdim = 0;
    uint32_t mindim = dim_num;
    for(int i = 0; i < dim_1;i++){
      auto index = *(axes->GetData() + i);
      if(index < mindim) mindim = index;
      if(index > maxdim) maxdim = index;
    }
    if(maxdim > dim_num){
      totalLength = context->GetInputTensor(0)->GetShapeSize();
      stepSize = 1;
      unitCount = 1;
    }else{
      for(int i=0;i<dim_num;i++){
        if(i < mindim){
          unitCount *= shape.GetDim(i);
        }else if(i>maxdim){
          stepSize *= shape.GetDim(i);
        }else{
          totalLength *= shape.GetDim(i);
        }
      }
    }
  }else{
    stepSize = 1;
    unitCount = 1;
    totalLength = context->GetInputTensor(0)->GetShapeSize();
  }
  tiling.set_totalLength(totalLength);
  tiling.set_stepSize(stepSize);
  tiling.set_unitCount(unitCount);
  printf("%u\t\t%u*****%u***%u*****\n", unitCount, totalLength, stepSize, pType);
  auto dt = context->GetInputTensor(0)->GetDataType();
  uint32_t typeKey = 0;
  uint32_t data_num = 0; // kernel中用到了多少组buffer和buf
  if (dt == 1) { // half
    sizeofdatatype = 2;
    data_num = 6;
    tiling.set_typeKey(1);
  }else if (dt == 0) { // float
    sizeofdatatype = 4;
    data_num = 3;
    tiling.set_typeKey(0);
  }

  // printf("%u***********\n", tiling.get_pType());
  uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
  //核内拆分，策略是尽可能的填满ub_size，最后一包单独处理，
  uint32_t ub_block_num =
      ((ub_size) / BLOCK_SIZE / data_num) * 13 / 20;  // ub_block在Ascend C中不能全部被用来作为输入输出，给了13/20系数。
  // uint32_t ub_block_num = 8;  //为测试方便，验证代码流程
  uint32_t tile_num;
  if(pType == 3){
    ub_block_num = 8;
  }
  if (ub_block_num % 2 != 0) {
    ub_block_num = ub_block_num - 1;
  }

  // 1.输入向量满足32字节对齐
  if (totalLength % ALIGN_NUM != 0) {  //不对齐，先32位对齐
    totalLengthAligned =
        ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
  } else {
    totalLengthAligned = totalLength;
  }

  
  context->SetBlockDim(1);

  uint32_t tileLength = 0;
  uint32_t lasttileLength = 0;

  tile_num = totalLengthAligned / ALIGN_NUM / ub_block_num; // 简单获取tile_num的大小
  if ((totalLengthAligned / ALIGN_NUM) % ub_block_num == 0 ||
      tile_num == 0) {  //对齐之后，可以均分到各个tile，或不足一个tile
    if (tile_num == 0) {
      tile_num = 1;
    } 
    if (totalLengthAligned < ub_block_num * ALIGN_NUM) {
      tileLength = totalLengthAligned;
      lasttileLength = tileLength;
    } else {
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = tileLength;
    }
  } else {  //满足字节对齐，核内不能均分
    tile_num = tile_num + 1;
    tileLength = ub_block_num * ALIGN_NUM;
    lasttileLength = totalLengthAligned - (tile_num - 1) * tileLength;
  }
  if(totalLength % tileLength){
    lasttileLength = totalLength % tileLength;
  }
  printf("%u\t\t%u******%u*******\n", tileLength, tile_num, lasttileLength);
  tiling.set_blockLength(totalLengthAligned);
  tiling.set_tileNum(tile_num);
  tiling.set_tileLength(tileLength);
  tiling.set_lasttileLength(lasttileLength);

  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t* currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;
  return ge::GRAPH_SUCCESS;
  
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    bool keepdim = context->GetAttrs()->GetBool(2);
    if(keepdim){
        *y_shape = *x1_shape;
    }else{
        gert::Shape out_shape({1});
        *y_shape = out_shape;
    }
    return GRAPH_SUCCESS;
}
}


namespace ops {
class LpNormV2 : public OpDef {
public:
    explicit LpNormV2(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("p").AttrType(OPTIONAL).Float(2);
        this->Attr("axes").AttrType(OPTIONAL).ListInt({});
        this->Attr("keepdim").AttrType(OPTIONAL).Bool(false);
        this->Attr("epsilon").AttrType(OPTIONAL).Float(1e-12);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(LpNormV2);
}
